import json
from typing import List
from pathlib import Path

import json_repair
from jinja2 import Environment
from research_agent.core.query import Query
from research_agent.core.general_llm import LLM
from pyaml_env import parse_config
from research_agent.core.config import Config
from jinja2 import Environment, FileSystemLoader
from research_agent.core.paths import UPDATE_REFERENCE_PROMPT


class UpdateReference:
    def __init__(self):
        configs = parse_config(Config.YAML_CONFIG)
        self.llm = LLM(config=configs[Config.DEFAULT_MODEL])

        self.prompt_env = Environment(
        loader=FileSystemLoader(UPDATE_REFERENCE_PROMPT))
        with open(UPDATE_REFERENCE_PROMPT, "r") as f:
            self.update_reference_prompt_template = self.prompt_env.from_string(
                f.read())

    async def update_reference(
        self,
        support_statement_citation_result: List[dict],
        paper_draft: str
    ) -> str:
        """
        参数说明：
        - support_statement_citation_result: 支持陈述和引用的字典列表
        - paper_draft: 需要更新的论文草稿
        返回更新后的论文草稿
        """
        # 生成提示信息
        prompt_messages = self._prepare_update_reference_prompt(
            support_statement_citation_result,
            paper_draft
        )
        # 调用 LLM

        response = await self.llm.completion(prompt_messages)
        # 返回更新后的草稿
        return response

    def _prepare_update_reference_prompt(self, support_data: List[dict], draft: str):
        """准备系统提示和用户提示的prompt组合"""
        system_prompt = self.update_reference_prompt_template.render(
            role="system")

        user_prompt = self.update_reference_prompt_template.render(
            role="user",
            support_statement_citation_result=support_data,
            paper_draft=draft
        )

        return [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]
